"""
多模态震动传播模型

实现了同时处理横向和纵向震动的多模态传播模型，包括：
- 横向和纵向震动的组合传播
- 多模态震动叠加
- 复杂环境下的震动传播模拟
"""

import numpy as np
from typing import Tuple, Dict, Optional
from loguru import logger
from .vibration_propagation import VibrationPropagation
from .decay_model import DecayModel


class MultiModalVibration:
    """多模态震动传播模型"""
    
    def __init__(self, 
                 horizontal_params: Optional[Dict] = None,
                 vertical_params: Optional[Dict] = None,
                 combination_method: str = "weighted_sum"):
        """
        初始化多模态震动传播模型
        
        Args:
            horizontal_params: 横向传播参数
            vertical_params: 纵向传播参数
            combination_method: 组合方法，支持weighted_sum, max, product
        """
        # 默认参数
        default_horizontal = {
            'propagation_speed': 1000.0,
            'step_size': 0.01,
            'attenuation_coefficient': 0.1
        }
        
        default_vertical = {
            'propagation_speed': 800.0,
            'step_size': 0.01,
            'attenuation_coefficient': 0.15
        }
        
        # 合并参数
        if horizontal_params:
            default_horizontal.update(horizontal_params)
        if vertical_params:
            default_vertical.update(vertical_params)
        
        # 初始化横向和纵向传播模型
        self.horizontal_propagation = VibrationPropagation(
            direction="horizontal", **default_horizontal
        )
        self.vertical_propagation = VibrationPropagation(
            direction="vertical", **default_vertical
        )
        
        self.combination_method = combination_method
        self.horizontal_weight = 0.6
        self.vertical_weight = 0.4
        
        logger.info("初始化多模态震动传播模型")
    
    def simulate_multi_modal_propagation(self, 
                                       initial_amplitude: float,
                                       source_position: Tuple[float, float],
                                       grid_size: Tuple[int, int] = (50, 50),
                                       grid_spacing: float = 0.01,
                                       time_steps: int = 100) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        模拟多模态震动传播
        
        Args:
            initial_amplitude: 初始振幅
            source_position: 震源位置 (x, y)
            grid_size: 网格大小 (rows, cols)
            grid_spacing: 网格间距 (m)
            time_steps: 时间步数
            
        Returns:
            (x坐标网格, y坐标网格, 组合振幅网格)
        """
        # 分别计算横向和纵向传播
        X_h, Y_h, amplitudes_h = self.horizontal_propagation.simulate_propagation(
            initial_amplitude, source_position, grid_size, grid_spacing, time_steps
        )
        
        X_v, Y_v, amplitudes_v = self.vertical_propagation.simulate_propagation(
            initial_amplitude, source_position, grid_size, grid_spacing, time_steps
        )
        
        # 组合多模态震动
        combined_amplitudes = self._combine_modalities(amplitudes_h, amplitudes_v)
        
        return X_h, Y_h, combined_amplitudes
    
    def _combine_modalities(self, 
                          horizontal_amplitudes: np.ndarray, 
                          vertical_amplitudes: np.ndarray) -> np.ndarray:
        """
        组合横向和纵向震动模态
        
        Args:
            horizontal_amplitudes: 横向振幅
            vertical_amplitudes: 纵向振幅
            
        Returns:
            组合后的振幅
        """
        if self.combination_method == "weighted_sum":
            return (self.horizontal_weight * horizontal_amplitudes + 
                   self.vertical_weight * vertical_amplitudes)
        
        elif self.combination_method == "max":
            return np.maximum(horizontal_amplitudes, vertical_amplitudes)
        
        elif self.combination_method == "product":
            return horizontal_amplitudes * vertical_amplitudes
        
        elif self.combination_method == "rms":
            return np.sqrt(horizontal_amplitudes**2 + vertical_amplitudes**2)
        
        else:
            logger.warning(f"未知的组合方法: {self.combination_method}，使用加权和")
            return (self.horizontal_weight * horizontal_amplitudes + 
                   self.vertical_weight * vertical_amplitudes)
    
    def propagate_to_targets(self, 
                           initial_amplitude: float,
                           source_position: Tuple[float, float],
                           target_positions: np.ndarray,
                           time_steps: int = 100) -> np.ndarray:
        """
        计算多模态震动传播到目标位置的振幅
        
        Args:
            initial_amplitude: 初始振幅
            source_position: 震源位置 (x, y)
            target_positions: 目标位置数组，形状为 (n, 2)
            time_steps: 时间步数
            
        Returns:
            目标位置的组合振幅数组，形状为 (n, time_steps)
        """
        # 分别计算横向和纵向传播
        horizontal_amplitudes = self.horizontal_propagation.propagate_vibration(
            initial_amplitude, source_position, target_positions, time_steps
        )
        
        vertical_amplitudes = self.vertical_propagation.propagate_vibration(
            initial_amplitude, source_position, target_positions, time_steps
        )
        
        # 组合多模态震动
        combined_amplitudes = self._combine_modalities(horizontal_amplitudes, vertical_amplitudes)
        
        return combined_amplitudes
    
    def plot_multi_modal_propagation(self, 
                                   initial_amplitude: float,
                                   source_position: Tuple[float, float],
                                   grid_size: Tuple[int, int] = (50, 50),
                                   grid_spacing: float = 0.01,
                                   title: str = None):
        """
        绘制多模态震动传播图
        
        Args:
            initial_amplitude: 初始振幅
            source_position: 震源位置
            grid_size: 网格大小
            grid_spacing: 网格间距
            title: 图表标题
        """
        try:
            import matplotlib.pyplot as plt
            
            X, Y, combined_amplitudes = self.simulate_multi_modal_propagation(
                initial_amplitude, source_position, grid_size, grid_spacing
            )
            
            # 创建子图
            fig, axes = plt.subplots(2, 2, figsize=(15, 12))
            
            # 横向传播
            X_h, Y_h, amplitudes_h = self.horizontal_propagation.simulate_propagation(
                initial_amplitude, source_position, grid_size, grid_spacing
            )
            contour1 = axes[0, 0].contourf(X_h, Y_h, amplitudes_h, levels=20, cmap='Blues')
            axes[0, 0].set_title('横向震动传播')
            axes[0, 0].set_xlabel('X 坐标 (m)')
            axes[0, 0].set_ylabel('Y 坐标 (m)')
            plt.colorbar(contour1, ax=axes[0, 0])
            
            # 纵向传播
            X_v, Y_v, amplitudes_v = self.vertical_propagation.simulate_propagation(
                initial_amplitude, source_position, grid_size, grid_spacing
            )
            contour2 = axes[0, 1].contourf(X_v, Y_v, amplitudes_v, levels=20, cmap='Reds')
            axes[0, 1].set_title('纵向震动传播')
            axes[0, 1].set_xlabel('X 坐标 (m)')
            axes[0, 1].set_ylabel('Y 坐标 (m)')
            plt.colorbar(contour2, ax=axes[0, 1])
            
            # 组合传播
            contour3 = axes[1, 0].contourf(X, Y, combined_amplitudes, levels=20, cmap='viridis')
            axes[1, 0].set_title(f'多模态震动传播 ({self.combination_method})')
            axes[1, 0].set_xlabel('X 坐标 (m)')
            axes[1, 0].set_ylabel('Y 坐标 (m)')
            plt.colorbar(contour3, ax=axes[1, 0])
            
            # 标记震源位置
            for ax in axes.flat:
                ax.plot(source_position[0], source_position[1], 'r*', 
                       markersize=15, label='震源')
                ax.grid(True, alpha=0.3)
                ax.axis('equal')
            
            # 振幅对比
            axes[1, 1].plot(amplitudes_h[25, :], 'b-', label='横向', linewidth=2)
            axes[1, 1].plot(amplitudes_v[25, :], 'r-', label='纵向', linewidth=2)
            axes[1, 1].plot(combined_amplitudes[25, :], 'g-', label='组合', linewidth=2)
            axes[1, 1].set_title('振幅对比 (中心线)')
            axes[1, 1].set_xlabel('Y 坐标')
            axes[1, 1].set_ylabel('振幅')
            axes[1, 1].legend()
            axes[1, 1].grid(True, alpha=0.3)
            
            plt.tight_layout()
            plt.suptitle(title or '多模态震动传播分析', y=1.02, fontsize=16)
            plt.show()
            
        except ImportError:
            logger.warning("matplotlib未安装，无法绘制多模态传播图")
    
    def set_combination_weights(self, horizontal_weight: float, vertical_weight: float):
        """
        设置组合权重
        
        Args:
            horizontal_weight: 横向权重
            vertical_weight: 纵向权重
        """
        total_weight = horizontal_weight + vertical_weight
        if total_weight != 0:
            self.horizontal_weight = horizontal_weight / total_weight
            self.vertical_weight = vertical_weight / total_weight
            logger.info(f"设置组合权重: 横向={self.horizontal_weight:.2f}, 纵向={self.vertical_weight:.2f}")
        else:
            logger.warning("权重总和不能为0")
    
    def get_model_parameters(self) -> Dict:
        """
        获取模型参数
        
        Returns:
            参数字典
        """
        return {
            'horizontal_params': self.horizontal_propagation.get_propagation_parameters(),
            'vertical_params': self.vertical_propagation.get_propagation_parameters(),
            'combination_method': self.combination_method,
            'horizontal_weight': self.horizontal_weight,
            'vertical_weight': self.vertical_weight
        }
    
    def optimize_model(self, 
                      experimental_data: np.ndarray,
                      target_positions: np.ndarray,
                      optimization_params: list = None) -> Dict:
        """
        根据实验数据优化模型参数
        
        Args:
            experimental_data: 实验数据
            target_positions: 目标位置
            optimization_params: 要优化的参数列表
            
        Returns:
            优化后的参数
        """
        # 这里可以实现参数优化算法
        # 例如：遗传算法、粒子群优化等
        logger.info("开始模型参数优化")
        
        # 简单的网格搜索优化示例
        best_params = {}
        best_error = float('inf')
        
        # 优化衰减系数
        for alpha_h in np.linspace(0.05, 0.2, 10):
            for alpha_v in np.linspace(0.1, 0.25, 10):
                self.horizontal_propagation.attenuation_coefficient = alpha_h
                self.vertical_propagation.attenuation_coefficient = alpha_v
                
                # 计算预测值
                predicted = self.propagate_to_targets(
                    1.0, (0, 0), target_positions
                )
                
                # 计算误差
                error = np.mean((predicted - experimental_data)**2)
                
                if error < best_error:
                    best_error = error
                    best_params = {
                        'horizontal_alpha': alpha_h,
                        'vertical_alpha': alpha_v,
                        'error': error
                    }
        
        logger.info(f"优化完成，最佳参数: {best_params}")
        return best_params
    
    def __str__(self) -> str:
        return (f"MultiModalVibration(combination={self.combination_method}, "
                f"h_weight={self.horizontal_weight:.2f}, v_weight={self.vertical_weight:.2f})")
    
    def __repr__(self) -> str:
        return self.__str__() 